from keras.layers import Input, Conv2D, Dropout, MaxPooling2D, Concatenate, UpSampling2D
from keras.optimizers import Adam, RMSprop
from keras.callbacks import ModelCheckpoint, TensorBoard, LambdaCallback
from keras.losses import binary_crossentropy,binary_focal_crossentropy
from keras.metrics import binary_accuracy
from keras.models import Model
from keras.regularizers import l1, l2

import cv2
import glob
import tensorflow as tf
import numpy as np

def unet(n_classes=2, input_shape=(255, 255, 3)):
    img_input = Input(shape=input_shape)
    
    conv1 = Conv2D(64, (3, 3), activation='relu', padding='same',kernel_initializer = 'he_normal')(img_input)
    conv1 = Conv2D(64, (3, 3), activation='relu', padding='same',kernel_initializer = 'he_normal')(conv1)
    pool1 = MaxPooling2D((2, 2), strides=2)(conv1)

    conv2 = Conv2D(128, (3, 3), activation='relu', padding='same',kernel_initializer = 'he_normal')(pool1)
    conv2 = Conv2D(128, (3, 3), activation='relu', padding='same',kernel_initializer = 'he_normal')(conv2)
    pool2 = MaxPooling2D((2, 2), strides=2)(conv2)

    conv3 = Conv2D(256, (3, 3), activation='relu', padding='same',kernel_initializer = 'he_normal')(pool2)
    conv3 = Conv2D(256, (3, 3), activation='relu', padding='same',kernel_initializer = 'he_normal')(conv3)
    pool3 = MaxPooling2D((2, 2), strides=2)(conv3)

    conv4 = Conv2D(512, (3, 3), activation='relu', padding='same',kernel_initializer = 'he_normal')(pool3)
    conv4 = Conv2D(512, (3, 3), activation='relu', padding='same',kernel_initializer = 'he_normal')(conv4)
    pool4 = MaxPooling2D((2, 2), strides=2)(conv4)
    
    conv5 = Conv2D(1024, (3, 3), activation='relu', padding='same',kernel_initializer = 'he_normal')(pool4)
    conv5 = Conv2D(1024, (3, 3), activation='relu', padding='same',kernel_initializer = 'he_normal')(conv5)
    conv5 = Dropout(0.5)(conv5)

    up6 = Conv2D(512, (3, 3), activation='relu', padding='same',kernel_initializer = 'he_normal')(UpSampling2D((2, 2))(conv5))
    up6 = tf.pad(up6,[[0,0],[0,conv4.shape[1]-up6.shape[1]],[0,conv4.shape[2]-up6.shape[2]],[0,0]])
    merge6 = Concatenate(axis=-1)([up6, conv4])
    conv6 = Conv2D(512, (3, 3), activation='relu', padding='same',kernel_initializer = 'he_normal')(merge6)
    conv6 = Conv2D(512, (3, 3), activation='relu', padding='same',kernel_initializer = 'he_normal')(conv6)

    up7 = Conv2D(256, (3, 3), activation='relu', padding='same',kernel_initializer = 'he_normal')(UpSampling2D((2, 2))(conv6))
    up7 = tf.pad(up7,[[0,0],[0,conv3.shape[1]-up7.shape[1]],[0,conv3.shape[2]-up7.shape[2]],[0,0]])
    merge7 = Concatenate(axis=-1)([up7, conv3])
    conv7 = Conv2D(256, (3, 3), activation='relu', padding='same',kernel_initializer = 'he_normal')(merge7)
    conv7 = Conv2D(256, (3, 3), activation='relu', padding='same',kernel_initializer = 'he_normal')(conv7)

    up8 = Conv2D(128, (3, 3), activation='relu', padding='same',kernel_initializer = 'he_normal')(UpSampling2D((2, 2))(conv7))
    up8 = tf.pad(up8,[[0,0],[0,conv2.shape[1]-up8.shape[1]],[0,conv2.shape[2]-up8.shape[2]],[0,0]])
    merge8 = Concatenate(axis=-1)([up8, conv2])
    conv8 = Conv2D(128, (3, 3), activation='relu', padding='same',kernel_initializer = 'he_normal')(merge8)
    conv8 = Conv2D(128, (3, 3), activation='relu', padding='same',kernel_initializer = 'he_normal')(conv8)

    up9 = Conv2D(64, (3, 3), activation='relu', padding='same',kernel_initializer = 'he_normal')(UpSampling2D((2, 2))(conv8))
    up9 = tf.pad(up9,[[0,0],[0,conv1.shape[1]-up9.shape[1]],[0,conv1.shape[2]-up9.shape[2]],[0,0]])
    merge9 = Concatenate(axis=-1)([up9, conv1])
    conv9 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge9)
    conv9 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv9)

    # 输出层
    outputs = Conv2D(n_classes, (1, 1), padding='same', activation='sigmoid' if n_classes==1 else 'softmax')(conv9)
    model = Model(img_input, outputs, name="unet")
    return model

from keras.preprocessing.image import ImageDataGenerator
def generator(image_directory,target_size,batch_size,dict_data,seed):
    # 定义图像生成器
    datagen_image = ImageDataGenerator(**dict_data)
    datagen_label = ImageDataGenerator(**dict_data)
    data_image = datagen_image.flow_from_directory(
            image_directory,
            classes=['images'],
            target_size=target_size,
            batch_size=batch_size,       #样本数不易过大
            seed=seed)
    data_label = datagen_label.flow_from_directory(
            image_directory,
            classes=["label"],
            color_mode='grayscale',
            target_size=target_size,
            batch_size=batch_size,       #样本数不易过大
            seed=seed)
    data = zip(data_image, data_label)
    for (i,img) in enumerate(data):
        cv2.imshow("image", img[0][0][0,:,:,:])
        cv2.imshow("label", img[1][0][0,:,:,:])
        cv2.waitKey(1)
        yield img[0][0],img[1][0]
    
if __name__ == "__main__":
    path = 'E:\\Opensource\\model\\mlc\\mlc_training_data\\'
    img = cv2.imread(glob.glob(path+"\\images\\*.png")[0])
    img = cv2.resize(img, (img.shape[1]//8*8, img.shape[0]//8*8))
    model = unet(1,img.shape)
    if True:
        model.summary()
        data = generator(path,img.shape[:2], 2, dict(
            rescale=1./255,rotation_range=10,width_shift_range=0.2,height_shift_range=0.2,
            shear_range=0.2,zoom_range=0.2,horizontal_flip=True,validation_split=0.2,fill_mode='constant'),1)
        model.compile(optimizer=Adam(learning_rate=1e-3), loss=binary_focal_crossentropy, metrics=['accuracy'])
        model.fit(data, steps_per_epoch=50, epochs=1, callbacks=
                [
                    TensorBoard('deep_net/logs',write_images=True,histogram_freq=1,write_graph=False),
                    ModelCheckpoint("deep_net\\param\\unet.keras",verbose=1),
                ])
    else:
        model.load_weights("deep_net\\param\\unet.keras")
        for file in glob.glob(path+"\\image\\*.png"):
            new_img = cv2.imread(file)/255.
            new_img = cv2.resize(new_img,(img.shape[1],img.shape[0]))
            new_img = new_img.reshape(1,*img.shape[:3])
            new_img = model.predict(new_img)
            cv2.imshow("",new_img[0,:,:,:])
            cv2.waitKey(1)
        